package ai.onnxruntime.example.objectdetection

import ai.onnxruntime.*
import ai.onnxruntime.extensions.OrtxPackage
import android.annotation.SuppressLint
import android.content.ContentValues
import android.content.pm.PackageManager
import android.graphics.Bitmap
import android.graphics.BitmapFactory
import android.graphics.Canvas
import android.graphics.Color
import android.graphics.Paint
import android.graphics.PorterDuff
import android.graphics.PorterDuffXfermode
import android.content.Intent
import android.os.Bundle
import android.provider.MediaStore
import android.util.Log
import android.view.View
import android.widget.Button
import android.widget.Toast
import androidx.activity.result.contract.ActivityResultContracts
import androidx.appcompat.app.AppCompatActivity
import androidx.camera.camera2.Camera2Config
import androidx.camera.core.*
import androidx.camera.lifecycle.ProcessCameraProvider
import androidx.camera.view.PreviewView
import androidx.core.content.ContextCompat
import com.github.chrisbanes.photoview.PhotoView
import java.io.ByteArrayInputStream
import java.io.ByteArrayOutputStream
import java.io.File
import java.io.InputStream
import java.text.SimpleDateFormat
import java.util.*
import java.util.concurrent.ExecutorService
import java.util.concurrent.Executors

class MainActivity : AppCompatActivity() {
    private var ortEnv: OrtEnvironment = OrtEnvironment.getEnvironment()
    private lateinit var ortSession: OrtSession
    private lateinit var inputPhotoView: PhotoView
    private lateinit var outputPhotoView: PhotoView
    private lateinit var previewView: PreviewView
    private lateinit var objectDetectionButton: Button
    private lateinit var selectFromGalleryButton: Button
    private lateinit var takePhotoButton: Button
    private lateinit var switchCameraButton: Button
    private lateinit var classes: List<String>
    private var currentBitmap: Bitmap? = null

    // CameraX相关变量
    private var imageCapture: ImageCapture? = null
    private var cameraProvider: ProcessCameraProvider? = null
    private var camera: Camera? = null
    private var lensFacing = CameraSelector.LENS_FACING_BACK
    private lateinit var cameraExecutor: ExecutorService

    // 权限请求
    private val requestPermissionLauncher = registerForActivityResult(
        ActivityResultContracts.RequestMultiplePermissions()
    ) { permissions ->
        val allGranted = permissions.values.all { it }
        if (allGranted) {
            startCamera()
            Toast.makeText(this, "权限已授予", Toast.LENGTH_SHORT).show()
        } else {
            Toast.makeText(this, "需要权限才能使用相机和相册功能", Toast.LENGTH_LONG).show()
        }
    }

    // 相册选择
    private val galleryResultLauncher = registerForActivityResult(
        ActivityResultContracts.StartActivityForResult()
    ) { result ->
        if (result.resultCode == RESULT_OK) {
            result.data?.data?.let { uri ->
                try {
                    val inputStream = contentResolver.openInputStream(uri)
                    currentBitmap = BitmapFactory.decodeStream(inputStream)
                    inputStream?.close()

                    currentBitmap?.let { bitmap ->
                        showPhotoView(bitmap)
                        Toast.makeText(this, "图片已选择", Toast.LENGTH_SHORT).show()
                    }
                } catch (e: Exception) {
                    Log.e(TAG, "Error loading image from gallery", e)
                    Toast.makeText(this, "加载图片失败", Toast.LENGTH_SHORT).show()
                }
            }
        }
    }

    @SuppressLint("UseCompatLoadingForDrawables")
    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        setContentView(R.layout.activity_main)

        // 初始化视图
        initViews()

        // 初始化模型和类
        classes = readClasses()
        initModel()

        // 初始化相机执行器
        cameraExecutor = Executors.newSingleThreadExecutor()

        // 设置按钮监听器
        setupButtonListeners()

        // 检查权限
        checkPermissions()

        previewView.post {
            // 确保预览视图可见
            previewView.visibility = View.VISIBLE
            inputPhotoView.visibility = View.GONE

            startCamera()
        }
    }

    private fun initViews() {
        inputPhotoView = findViewById(R.id.photoView1)
        outputPhotoView = findViewById(R.id.photoView2)
        previewView = findViewById(R.id.previewView)
        objectDetectionButton = findViewById(R.id.object_detection_button)
        selectFromGalleryButton = findViewById(R.id.select_from_gallery_button)
        takePhotoButton = findViewById(R.id.take_photo_button)
        switchCameraButton = findViewById(R.id.switch_camera_button)
    }

    private fun initModel() {
        val sessionOptions: OrtSession.SessionOptions = OrtSession.SessionOptions()
        sessionOptions.registerCustomOpLibrary(OrtxPackage.getLibraryPath())

        try {
            ortSession = ortEnv.createSession(readModel(), sessionOptions)
        } catch (e: Exception) {
            Log.e(TAG, "Failed to create ORT session", e)
            Toast.makeText(this, "模型加载失败", Toast.LENGTH_LONG).show()
        }
    }

    private fun checkPermissions() {
        val requiredPermissions = arrayOf(
            android.Manifest.permission.CAMERA,
            android.Manifest.permission.READ_EXTERNAL_STORAGE,
            android.Manifest.permission.WRITE_EXTERNAL_STORAGE
        )

        val missingPermissions = requiredPermissions.filter {
            ContextCompat.checkSelfPermission(this, it) != PackageManager.PERMISSION_GRANTED
        }

        if (missingPermissions.isNotEmpty()) {
            requestPermissionLauncher.launch(missingPermissions.toTypedArray())
        } else {
            startCamera()
        }
    }

    private fun setupButtonListeners() {
        selectFromGalleryButton.setOnClickListener {
            val intent = Intent(Intent.ACTION_PICK, MediaStore.Images.Media.EXTERNAL_CONTENT_URI)
            galleryResultLauncher.launch(intent)
        }

        takePhotoButton.setOnClickListener {
            takePhoto()
        }

        switchCameraButton.setOnClickListener {
            switchCamera()
        }

        objectDetectionButton.setOnClickListener {
            currentBitmap?.let { bitmap ->
                Thread {
                    try {
                        val result = performObjectDetection(bitmap)
                        runOnUiThread {
                            updateUI(result)
                            Toast.makeText(baseContext, "物体检测完成!", Toast.LENGTH_SHORT).show()
                        }
                    } catch (e: Exception) {
                        Log.e(TAG, "Exception caught when perform ObjectDetection", e)
                        runOnUiThread {
                            Toast.makeText(baseContext, "物体检测失败", Toast.LENGTH_SHORT).show()
                        }
                    }
                }.start()
            } ?: run {
                Toast.makeText(baseContext, "请先选择图片", Toast.LENGTH_SHORT).show()
            }
        }
    }

    // CameraX相机启动
    private fun startCamera() {
        val cameraProviderFuture = ProcessCameraProvider.getInstance(this)

        cameraProviderFuture.addListener({
            try {
                cameraProvider = cameraProviderFuture.get()

                // 检查设备是否有摄像头
                if (cameraProvider?.hasCamera(CameraSelector.DEFAULT_BACK_CAMERA) != true &&
                    cameraProvider?.hasCamera(CameraSelector.DEFAULT_FRONT_CAMERA) != true) {
                    Toast.makeText(this, "设备没有可用的摄像头", Toast.LENGTH_LONG).show()
                    return@addListener
                }

                // 创建预览用例
                val preview = Preview.Builder()
                    .build()
                    .also {
                        it.setSurfaceProvider(previewView.surfaceProvider)
                    }

                // 创建拍照用例
                imageCapture = ImageCapture.Builder()
                    .setCaptureMode(ImageCapture.CAPTURE_MODE_MINIMIZE_LATENCY)
                    .build()

                // 确保选择的镜头方向可用
                val availableCameras = cameraProvider?.availableCameraInfos ?: emptyList()
                val cameraSelector = if (availableCameras.isNotEmpty()) {
                    CameraSelector.Builder()
                        .requireLensFacing(lensFacing)
                        .build()
                } else {
                    CameraSelector.DEFAULT_BACK_CAMERA
                }

                // 解绑所有用例
                cameraProvider?.unbindAll()

                // 绑定用例到生命周期
                camera = cameraProvider?.bindToLifecycle(
                    this, cameraSelector, preview, imageCapture
                )

                Log.d(TAG, "相机启动成功，镜头方向: $lensFacing")

            } catch (exc: Exception) {
                Log.e(TAG, "Use case binding failed", exc)
                Toast.makeText(this, "相机启动失败: ${exc.message}", Toast.LENGTH_LONG).show()
            }

        }, ContextCompat.getMainExecutor(this))
    }

    // 拍照功能
    private fun takePhoto() {
        val imageCapture = imageCapture ?: return

        // 照片存本地
//        // 创建时间戳名称
//        val name = SimpleDateFormat("yyyy-MM-dd-HH-mm-ss-SSS", Locale.US)
//            .format(System.currentTimeMillis())
//
//        val contentValues = ContentValues().apply {
//            put(MediaStore.Images.Media.DISPLAY_NAME, name)
//            put(MediaStore.Images.Media.MIME_TYPE, "image/jpeg")
//            if (android.os.Build.VERSION.SDK_INT > android.os.Build.VERSION_CODES.P) {
//                put(MediaStore.Images.Media.RELATIVE_PATH, "Pictures/ObjectDetection")
//            }
//        }
//
//        // 创建输出选项
//        val outputOptions = ImageCapture.OutputFileOptions.Builder(
//            contentResolver,
//            MediaStore.Images.Media.EXTERNAL_CONTENT_URI,
//            contentValues
//        ).build

        // 创建临时文件，放在应用缓存目录
        val tempFile = File.createTempFile(
            "temp_capture_${System.currentTimeMillis()}",
            ".jpg",
            cacheDir // 使用缓存目录，系统可自动清理
        )

        val outputOptions = ImageCapture.OutputFileOptions.Builder(tempFile).build()


        // 执行拍照
        imageCapture.takePicture(
            outputOptions,
            ContextCompat.getMainExecutor(this),
            object : ImageCapture.OnImageSavedCallback {
                override fun onImageSaved(output: ImageCapture.OutputFileResults) {
//                    val savedUri = output.savedUri ?: return
//
//                    // 加载拍摄的图片
//                    try {
//                        val inputStream = contentResolver.openInputStream(savedUri)
//                        currentBitmap = BitmapFactory.decodeStream(inputStream)
//                        inputStream?.close()
//
//                        currentBitmap?.let { bitmap ->
//                            showPhotoView(bitmap)
//                            Toast.makeText(baseContext, "照片已保存", Toast.LENGTH_SHORT).show()
//                        }
//                    } catch (e: Exception) {
//                        Log.e(TAG, "Error loading captured photo", e)
//                    }
                    try {
                        // 从临时文件读取 Bitmap
                        currentBitmap = BitmapFactory.decodeFile(tempFile.absolutePath)

                        runOnUiThread {
                            currentBitmap?.let {
                                showPhotoView(it)
                                Toast.makeText(this@MainActivity, "照片已捕获", Toast.LENGTH_SHORT).show()
                            }
                        }
                    } finally {
                        // 关键：无论成功与否，都删除临时文件
                        tempFile.delete()
                    }
                }
                override fun onError(exc: ImageCaptureException) {
                    Log.e(TAG, "Photo capture failed: ${exc.message}", exc)
                    // 确保出错时也删除临时文件
                    tempFile.delete()
                    Toast.makeText(baseContext, "拍照失败", Toast.LENGTH_SHORT).show()
                }
            }
        )
    }

    // 切换摄像头
    private fun switchCamera() {
        lensFacing = if (lensFacing == CameraSelector.LENS_FACING_BACK) {
            CameraSelector.LENS_FACING_FRONT
        } else {
            CameraSelector.LENS_FACING_BACK
        }
        showCameraPreview()
        startCamera()
    }

    // 显示图片视图（隐藏预览）
    private fun showPhotoView(bitmap: Bitmap) {
        inputPhotoView.setImageBitmap(bitmap)
        previewView.visibility = View.GONE
        inputPhotoView.visibility = View.VISIBLE
    }

    // 显示相机预览（隐藏图片视图）
    private fun showCameraPreview() {
        previewView.visibility = View.VISIBLE
        inputPhotoView.visibility = View.GONE
    }

    // 以下方法保持原有功能不变
    private fun updateUI(result: Result) {
        val mutableBitmap: Bitmap = result.outputBitmap.copy(Bitmap.Config.ARGB_8888, true)
        val canvas = Canvas(mutableBitmap)
        val paint = Paint().apply {
            color = Color.WHITE
            textSize = 28f
            xfermode = PorterDuffXfermode(PorterDuff.Mode.SRC_OVER)
        }

        canvas.drawBitmap(mutableBitmap, 0.0f, 0.0f, paint)

        result.outputBox.iterator().forEach { boxInfo ->
            canvas.drawText(
                "%s:%.2f".format(classes[boxInfo[5].toInt()], boxInfo[4]),
                boxInfo[0] - boxInfo[2] / 2,
                boxInfo[1] - boxInfo[3] / 2,
                paint
            )
        }

        outputPhotoView.setImageBitmap(mutableBitmap)
    }

    private fun performObjectDetection(bitmap: Bitmap): Result {
        val inputStream = bitmapToInputStream(bitmap)
        val objDetector = ObjectDetector()
        return objDetector.detect(inputStream, ortEnv, ortSession)
    }

    private fun bitmapToInputStream(bitmap: Bitmap): InputStream {
        val stream = ByteArrayOutputStream()
        bitmap.compress(Bitmap.CompressFormat.JPEG, 100, stream)
        return ByteArrayInputStream(stream.toByteArray())
    }

    private fun readModel(): ByteArray {
        val modelID = R.raw.yolov8n_with_pre_post_processing
        return resources.openRawResource(modelID).readBytes()
    }

    private fun readClasses(): List<String> {
        return resources.openRawResource(R.raw.classes).bufferedReader().readLines()
    }

    override fun onDestroy() {
        super.onDestroy()
        cameraExecutor.shutdown()
        ortSession.close()
        ortEnv.close()
    }

    companion object {
        const val TAG = "ORTObjectDetection"
    }
}